# Copyright (c) 2020 Presto Labs Pte. Ltd.
# Author: chunhui

import numpy as np
import pandas as pd

from coin.exchange.base.order_gateway import OrderSide
from ..on_log import OnLog
from .base import ExportDataFrame, OrderInfoTracker


class FillPositionMismatch(OrderInfoTracker, ExportDataFrame, OnLog):
  def __init__(self, tolerate_interval=5, epsilon=1e-3):
    super().__init__()
    self.epsilon = epsilon
    self.tolerate_interval = tolerate_interval
    self.pos_by_query = {}
    self.pos_by_fill = {}
    self.diff_start_ts = {}
    self.last_order_info = {}
    self.pos_range = {}
    self.excluded_ccy = None
    self.is_fut = None
    self.counter = 0

  @staticmethod
  def parse_futures_info(arg, market_type, ccy):
    handler = eval(arg)
    handler._update_futures_info(market_type == 'Futures', ccy)
    return handler

  def _update_futures_info(self, is_futures, ccy):
    self.is_fut = is_futures
    if not self.is_fut:
      self.excluded_ccy = ccy

  def _get_next_dummy_order_id(self):
    self.counter += 1
    return "dummy" + str(self.counter)

  def _mark_errors(self, ts, order_info, currency):
    order_id = order_info[0] if order_info else self._get_next_dummy_order_id()
    fill_ts = order_info[2] if order_info else ''
    self.errors[order_id] = {
        "order_id": order_id if order_info else '',
        "log_ts": ts,
        "currency": currency,
        "query_pos": round(self.pos_by_query[currency], 3),
        "fill_pos": round(self.pos_by_fill[currency], 3),
        "diff": round(abs(self.pos_by_fill[currency] - self.pos_by_query[currency]), 3),
    }

    if order_info and order_id in self.order_info:
      fill_qty = order_info[1]
      self.errors[order_id].update(self.order_info[order_id])
      del self.errors[order_id]['qty']
      extra_dict = {"order_life": round((fill_ts - self.order_info[order_id]['submit_ts'])
                                        / 1e9, 3), "fill_qty": fill_qty}
      self.errors[order_id].update(extra_dict)
    else:
      extra_fields = ['fill_qty', 'order_life', 'side', 'symbol', 'price', "submit_ts"]
      self.errors[order_id].update({x: '' for x in extra_fields})
    if currency in self.pos_range:
      diff_info = {"pos_range": self.pos_range[currency],
                   "relative_diff": round(abs((self.pos_by_fill[currency]
                                               - self.pos_by_query[currency])
                                               / (self.pos_range[currency]+1e-20)) * 100, 3)}
    else:
      diff_info = {"pos_range": np.nan, "relative_diff": np.nan}

    self.errors[order_id].update(diff_info)
    self.errors[order_id]['submit_ts'] = fill_ts

  def on_config(self, ts, symbol, lot_size, min_pos, max_pos):
    ccy = symbol if self.is_fut else symbol.split('-')[0]
    self.pos_range[ccy] = max_pos - min_pos

  def on_balance(self, ts, currency, balance):
    if self.is_fut or currency == self.excluded_ccy:
      return

    if currency not in self.pos_by_fill:  # Initialize
      self.pos_by_fill[currency] = balance
      self.pos_by_query[currency] = balance
      self.diff_start_ts[currency] = None

    if abs(self.pos_by_query[currency] - self.pos_by_fill[currency]) > self.epsilon \
          and self.diff_start_ts.get(currency, None) is not None \
          and int((ts - self.diff_start_ts[currency]) / 1e9) > self.tolerate_interval:
      self._mark_errors(ts, self.last_order_info[currency], currency)
      self.pos_by_query[currency] = balance
      self.pos_by_fill[currency] = self.pos_by_query[currency]
      self.diff_start_ts[currency] = None
      return

    self.pos_by_query[currency] = balance
    if abs(balance - self.pos_by_fill[currency]) < self.epsilon:
      self.diff_start_ts[currency] = None
    else:
      if self.diff_start_ts[currency] is None:
        self.diff_start_ts[currency] = ts
        self.last_order_info[currency] = None

  def on_position(self, ts, currency, position):
    if not self.is_fut:
      return
    if currency not in self.pos_by_fill:  # Initialize
      self.pos_by_fill[currency] = position
      self.pos_by_query[currency] = position
      self.diff_start_ts[currency] = None

    if abs(self.pos_by_query[currency] - self.pos_by_fill[currency]) > self.epsilon \
          and self.diff_start_ts.get(currency, None) is not None \
          and int((ts - self.diff_start_ts[currency]) / 1e9) > self.tolerate_interval:
      self._mark_errors(ts, self.last_order_info[currency], currency)
      self.pos_by_query[currency] = position
      self.pos_by_fill[currency] = self.pos_by_query[currency]
      self.diff_start_ts[currency] = None
      return

    self.pos_by_query[currency] = position
    if abs(position - self.pos_by_fill[currency]) < self.epsilon:
      self.diff_start_ts[currency] = None
    else:
      if self.diff_start_ts[currency] is None:
        self.diff_start_ts[currency] = ts
        self.last_order_info[currency] = None

  def on_fill(self, ts, currency, order_id, order_side, fill_price, fill_qty, maker):
    qty = fill_qty if order_side == OrderSide.BUY else -fill_qty
    ccy = currency if self.is_fut else currency.split('-')[0]
    if ccy not in self.pos_by_fill:  # Uninitialized
      return
    if abs(self.pos_by_query[ccy] - self.pos_by_fill[ccy]) > self.epsilon \
          and self.diff_start_ts[ccy] is not None \
          and int((ts - self.diff_start_ts[ccy]) / 1e9) > self.tolerate_interval:
      self._mark_errors(ts, self.last_order_info[ccy], ccy)
      self.pos_by_fill[ccy] = self.pos_by_query[ccy]
      self.diff_start_ts[ccy] = None
      return

    self.pos_by_fill[ccy] += qty
    if abs(self.pos_by_query[ccy] - self.pos_by_fill[ccy]) < self.epsilon:
      self.diff_start_ts[ccy] = None
    else:
      if self.diff_start_ts[ccy] is None or self.last_order_info[ccy] is None:
        # Avoid high frequency fill msg overriding the starting time
        self.diff_start_ts[ccy] = ts
        self.last_order_info[ccy] = (order_id, fill_qty, ts)

  def to_df(self):
    self.ts_fields.append('log_ts')
    df = super().to_df().sort_values('relative_diff', ascending=False).reset_index()
    df = df.rename(columns={'submit_ts': 'fill_ts'})
    df['relative_diff'] = df['relative_diff'].astype(str) + "%"
    df.loc[np.isnan(df['pos_range']), 'relative_diff'] = ''
    df.loc[np.isnan(df['pos_range']), 'pos_range'] = ''
    df.loc[pd.isnull(df['fill_ts']), 'fill_ts'] = ''
    cols = ['log_ts', 'currency', 'query_pos', 'fill_pos', 'diff', 'relative_diff', 'pos_range',
            'order_id', 'order_life', 'symbol', 'price', 'fill_qty', 'side', 'fill_ts']
    df = df[cols].set_index('log_ts')
    return df
